#include "viml/lexer.h"
#include <stdexcept>
#include <cctype>

namespace viml {

// 关键字映射表初始化
std::map<std::string, TokenType> Lexer::keywords = {
    {"let", TokenType::LET},
    {"if", TokenType::IF},
    {"elseif", TokenType::ELSEIF},
    {"else", TokenType::ELSE},
    {"endif", TokenType::ENDIF},
    {"while", TokenType::WHILE},
    {"endwhile", TokenType::ENDWHILE},
    {"for", TokenType::FOR},
    {"in", TokenType::IN},
    {"endfor", TokenType::ENDFOR},
    {"function", TokenType::FUNCTION},
    {"endfunction", TokenType::ENDFUNCTION},
    {"return", TokenType::RETURN},
    {"echo", TokenType::ECHO},
    {"call", TokenType::CALL}
};

Lexer::Lexer(const std::string& src) 
    : source(src), position(0), line(1), column(1), current_char('\0') {
    if (!source.empty()) {
        current_char = source[0];
    }
}

void Lexer::advance() {
    if (position < source.length() - 1) {
        position++;
        column++;
        if (current_char == '\n') {
            line++;
            column = 1;
        }
        current_char = source[position];
    } else {
        current_char = '\0';
    }
}

char Lexer::peek() {
    size_t peek_pos = position + 1;
    if (peek_pos < source.length()) {
        return source[peek_pos];
    }
    return '\0';
}

void Lexer::skip_whitespace() {
    while (current_char != '\0' && std::isspace(current_char) && current_char != '\n') {
        advance();
    }
}

void Lexer::skip_comment() {
    // VimL注释以 " 开始（只有当不在字符串上下文中时）
    if (current_char == '"') {
        // 检查是否为行注释（" 必须是行的第一个非空白字符或前面有空格）
        advance(); // 跳过 "
        while (current_char != '\0' && current_char != '\n') {
            advance();
        }
    }
}

std::string Lexer::read_number() {
    std::string result;
    
    while (current_char != '\0' && (std::isdigit(current_char) || current_char == '.')) {
        result += current_char;
        advance();
    }
    
    return result;
}

std::string Lexer::read_string(char quote_type) {
    std::string result;
    char quote = (quote_type == '\0') ? current_char : quote_type; // ' 或 "
    advance(); // 跳过开始引号
    
    while (current_char != '\0' && current_char != quote) {
        if (current_char == '\\') {
            advance();
            if (current_char != '\0') {
                switch (current_char) {
                    case 'n': result += '\n'; break;
                    case 't': result += '\t'; break;
                    case 'r': result += '\r'; break;
                    case '\\': result += '\\'; break;
                    case '\'': result += '\''; break;
                    case '"': result += '"'; break;
                    default: result += current_char; break;
                }
                advance();
            }
        } else {
            result += current_char;
            advance();
        }
    }
    
    if (current_char == quote) {
        advance(); // 跳过结束引号
    } else {
        error("Unterminated string literal");
    }
    
    return result;
}

std::string Lexer::read_identifier() {
    std::string result;
    
    while (current_char != '\0' && 
           (std::isalnum(current_char) || current_char == '_' || current_char == ':' || current_char == '$')) {
        result += current_char;
        advance();
    }
    
    return result;
}

TokenType Lexer::get_keyword_type(const std::string& word) {
    auto it = keywords.find(word);
    if (it != keywords.end()) {
        return it->second;
    }
    return TokenType::IDENTIFIER;
}

Token Lexer::next_token() {
    while (current_char != '\0') {
        size_t token_line = line;
        size_t token_column = column;
        
        // 跳过空白字符（除了换行符）
        if (std::isspace(current_char) && current_char != '\n') {
            skip_whitespace();
            continue;
        }
        
        // 处理双引号：可能是字符串字面量或行尾注释
        if (current_char == '"') {
            // 判断是字符串还是注释
            bool is_comment = false;
            
            // 如果在行首或者前面是空白字符，则可能是注释或字符串
            if (token_column == 1 || (position > 0 && std::isspace(source[position - 1]))) {
                // 查找行尾或匹配的引号
                size_t search_pos = position + 1;
                bool found_closing_quote = false;
                
                // 查找同一行内是否有匹配的引号
                while (search_pos < source.length() && source[search_pos] != '\n') {
                    if (source[search_pos] == '"') {
                        found_closing_quote = true;
                        break;
                    }
                    search_pos++;
                }
                
                // 如果没有找到匹配的引号，则视为注释
                if (!found_closing_quote) {
                    is_comment = true;
                }
            }
            
            if (is_comment) {
                // 处理行尾注释
                skip_comment();
                continue;
            } else {
                // 处理字符串字面量
                return Token(TokenType::STRING, read_string('"'), token_line, token_column);
            }
        }
        
        // 字符串（单引号）
        if (current_char == '\'') {
            return Token(TokenType::STRING, read_string('\''), token_line, token_column);
        }
        
        // 换行符
        if (current_char == '\n') {
            advance();
            return Token(TokenType::NEWLINE, "\\n", token_line, token_column);
        }
        
        // 数字
        if (std::isdigit(current_char)) {
            return Token(TokenType::NUMBER, read_number(), token_line, token_column);
        }
        
        // 特殊常量 (v:开头)
        if (current_char == 'v' && peek() == ':') {
            std::string word = read_identifier();
            return Token(TokenType::SPECIAL_CONSTANT, word, token_line, token_column);
        }
        
        // 环境变量 ($开头)
        if (current_char == '$') {
            std::string word = read_identifier();
            return Token(TokenType::IDENTIFIER, word, token_line, token_column);
        }
        
        // 标识符和关键字
        if (std::isalpha(current_char) || current_char == '_') {
            std::string word = read_identifier();
            TokenType type = get_keyword_type(word);
            return Token(type, word, token_line, token_column);
        }
        
        // 双字符操作符
        if (current_char == '=' && peek() == '=') {
            advance();
            advance();
            return Token(TokenType::EQUAL, "==", token_line, token_column);
        }
        
        if (current_char == '!' && peek() == '=') {
            advance();
            advance();
            return Token(TokenType::NOT_EQUAL, "!=", token_line, token_column);
        }
        
        if (current_char == '<' && peek() == '=') {
            advance();
            advance();
            return Token(TokenType::LESS_EQUAL, "<=", token_line, token_column);
        }
        
        if (current_char == '>' && peek() == '=') {
            advance();
            advance();
            return Token(TokenType::GREATER_EQUAL, ">=", token_line, token_column);
        }
        
        if (current_char == '&' && peek() == '&') {
            advance();
            advance();
            return Token(TokenType::AND, "&&", token_line, token_column);
        }
        
        if (current_char == '|' && peek() == '|') {
            advance();
            advance();
            return Token(TokenType::OR, "||", token_line, token_column);
        }
        
        // 单字符操作符和分隔符
        char ch = current_char;
        advance();
        
        switch (ch) {
            case '=': return Token(TokenType::ASSIGN, "=", token_line, token_column);
            case '+': return Token(TokenType::PLUS, "+", token_line, token_column);
            case '-': return Token(TokenType::MINUS, "-", token_line, token_column);
            case '*': return Token(TokenType::MULTIPLY, "*", token_line, token_column);
            case '/': return Token(TokenType::DIVIDE, "/", token_line, token_column);
            case '%': return Token(TokenType::MODULO, "%", token_line, token_column);
            case '<': return Token(TokenType::LESS_THAN, "<", token_line, token_column);
            case '>': return Token(TokenType::GREATER_THAN, ">", token_line, token_column);
            case '!': return Token(TokenType::NOT, "!", token_line, token_column);
            case '(': return Token(TokenType::LPAREN, "(", token_line, token_column);
            case ')': return Token(TokenType::RPAREN, ")", token_line, token_column);
            case '[': return Token(TokenType::LBRACKET, "[", token_line, token_column);
            case ']': return Token(TokenType::RBRACKET, "]", token_line, token_column);
            case '{': return Token(TokenType::LBRACE, "{", token_line, token_column);
            case '}': return Token(TokenType::RBRACE, "}", token_line, token_column);
            case ',': return Token(TokenType::COMMA, ",", token_line, token_column);
            case '.': return Token(TokenType::DOT, ".", token_line, token_column);
            case ':': return Token(TokenType::COLON, ":", token_line, token_column);
            case ';': return Token(TokenType::SEMICOLON, ";", token_line, token_column);
            default:
                return Token(TokenType::UNKNOWN, std::string(1, ch), token_line, token_column);
        }
    }
    
    return Token(TokenType::END_OF_FILE, "", line, column);
}

Token Lexer::peek_token() {
    // 保存当前状态
    size_t saved_position = position;
    size_t saved_line = line;
    size_t saved_column = column;
    char saved_char = current_char;
    
    // 获取下一个token
    Token token = next_token();
    
    // 恢复状态
    position = saved_position;
    line = saved_line;
    column = saved_column;
    current_char = saved_char;
    
    return token;
}

std::vector<Token> Lexer::tokenize() {
    std::vector<Token> tokens;
    
    Token token = next_token();
    while (token.type != TokenType::END_OF_FILE) {
        tokens.push_back(token);
        token = next_token();
    }
    tokens.push_back(token); // 添加EOF token
    
    return tokens;
}

void Lexer::error(const std::string& message) {
    throw std::runtime_error("Lexer error at line " + std::to_string(line) + 
                            ", column " + std::to_string(column) + ": " + message);
}

// Token类型转换为字符串
std::string token_type_to_string(TokenType type) {
    switch (type) {
        case TokenType::NUMBER: return "NUMBER";
        case TokenType::STRING: return "STRING";
        case TokenType::IDENTIFIER: return "IDENTIFIER";
        case TokenType::LET: return "LET";
        case TokenType::IF: return "IF";
        case TokenType::ELSEIF: return "ELSEIF";
        case TokenType::ELSE: return "ELSE";
        case TokenType::ENDIF: return "ENDIF";
        case TokenType::WHILE: return "WHILE";
        case TokenType::ENDWHILE: return "ENDWHILE";
        case TokenType::FOR: return "FOR";
        case TokenType::IN: return "IN";
        case TokenType::ENDFOR: return "ENDFOR";
        case TokenType::FUNCTION: return "FUNCTION";
        case TokenType::ENDFUNCTION: return "ENDFUNCTION";
        case TokenType::RETURN: return "RETURN";
        case TokenType::ECHO: return "ECHO";
        case TokenType::CALL: return "CALL";
        case TokenType::ASSIGN: return "ASSIGN";
        case TokenType::PLUS: return "PLUS";
        case TokenType::MINUS: return "MINUS";
        case TokenType::MULTIPLY: return "MULTIPLY";
        case TokenType::DIVIDE: return "DIVIDE";
        case TokenType::MODULO: return "MODULO";
        case TokenType::EQUAL: return "EQUAL";
        case TokenType::NOT_EQUAL: return "NOT_EQUAL";
        case TokenType::LESS_THAN: return "LESS_THAN";
        case TokenType::LESS_EQUAL: return "LESS_EQUAL";
        case TokenType::GREATER_THAN: return "GREATER_THAN";
        case TokenType::GREATER_EQUAL: return "GREATER_EQUAL";
        case TokenType::AND: return "AND";
        case TokenType::OR: return "OR";
        case TokenType::NOT: return "NOT";
        case TokenType::LPAREN: return "LPAREN";
        case TokenType::RPAREN: return "RPAREN";
        case TokenType::LBRACKET: return "LBRACKET";
        case TokenType::RBRACKET: return "RBRACKET";
        case TokenType::LBRACE: return "LBRACE";
        case TokenType::RBRACE: return "RBRACE";
        case TokenType::COMMA: return "COMMA";
        case TokenType::DOT: return "DOT";
        case TokenType::COLON: return "COLON";
        case TokenType::SEMICOLON: return "SEMICOLON";
        case TokenType::NEWLINE: return "NEWLINE";
        case TokenType::SPECIAL_CONSTANT: return "SPECIAL_CONSTANT";
        case TokenType::END_OF_FILE: return "END_OF_FILE";
        case TokenType::UNKNOWN: return "UNKNOWN";
        default: return "UNKNOWN";
    }
}

} // namespace viml